-- | Module:        Hexena.Chunks
--   Stability:     Experimental
--   Portability:   Portable
--  
-- <TODO>
module Hexena.Chunks
  ( Chunkable
  , IFFMap
  , Testtype
  , bodyID
  , isForm
  , isFormId
  , loadBODY
  , loadMalwIFF
  , loadMalwIFF'
  , loadMalwareFile
  , loadANNO
  , saveANNO
  -- , saveBLEN
  , loadFORMHIST
  , makeFORMHIST
  , mapIFF
  , loadPEHD
  , savePEHD
  , saveHASH
  , malwId
  , writeIFF
  ) where

import qualified Codec.Compression.BZip as CBZ
import qualified Codec.Compression.GZip as CGZ
import Control.Monad (ap)
import Data.Array.Unboxed (elems, listArray)
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Char (chr, ord)
import qualified Data.Digest.SHA1 as SHA1
import Data.IFF
import Data.List (find)
import qualified Data.Map as M
import Data.Maybe

import Hexena.PEHeader
import Hexena.Math.Statistics

data Testtype = Testtype Word32 Word16
type IFFMap = M.Map String Chunk

bodyID :: ChunkId
bodyID = chunkIdFromString "BODY"

isFormId :: ChunkId -> T -> Bool
isFormId cid (Cons c (Form _)) = c == cid
isFormId _ _ = False

isForm :: T -> Bool
isForm (Cons _ (Form _)) = True
isForm _ = False

-- | Load a malware file and return either a hexena file of just a
-- bytestring.
loadMalwareFile :: String -> IO (Either BL.ByteString T)
loadMalwareFile fn = let pIFF Nothing = error "File is not an IFF file!"
                         pIFF (Just x) = x
                     in do x <- B.readFile fn
                           if [0x4d, 0x5a] == take 2 (B.unpack x) then return (Left $ BL.fromChunks [x]) else return (Right $ pIFF $ fromByteString x)

-- | Load a MALW IFF file. It will throw an exception if the file can
-- | not be loaded, is not an IFF file or if it is not a FORM.MALW IFF
-- | file.
loadMalwIFF :: String -> IO IFFMap
loadMalwIFF fname = do dat <- fromByteString `fmap` B.readFile fname
                       case dat of
                         Nothing -> error "File is not an IFF file!"
                         Just x -> return $ mapIFF x

-- | Load a MALW IFF file. It will throw an exception if the file can
-- | not be loaded, is not an IFF file or if it is not a FORM.MALW IFF
-- | file.
loadMalwIFF' :: String -> IO (IFFMap, [T])
loadMalwIFF' fname = let deiff (Cons cid (Form xs)) = xs
                     in do dat <- fromByteString `fmap` B.readFile fname
                           case dat of
                             Nothing -> error "File is not an IFF file!"
                             Just x -> return (mapIFF x, deiff x)


mapIFF :: T -> IFFMap
mapIFF (Cons cid (Form xs)) | cid == malwId = M.fromList $ map f xs 
                            | otherwise = error "Guru Meditation! File is not an MALW IFF file!"
    where f (Cons i c) = (chunkIdToString i, c)
mapIFF _ = error "Guru Meditation! File is not a FORM IFF file!"

decompressBodyF :: Word8 -> B.ByteString -> B.ByteString
decompressBodyF 0 = id
decompressBodyF 2 = \x -> B.concat $ BL.toChunks $ CGZ.decompress $ BL.fromChunks [x]
decompressBodyF 3 = \x -> B.concat $ BL.toChunks $ CBZ.decompress $ BL.fromChunks [x]
decompressBodyF _ = error "Illegal compression level"

loadBODY :: [T] -> B.ByteString
loadBODY iff = let body = content $ chunk $ fromJust $ find ((bodyID ==) . chunkId) iff
                   mwhd = content $ chunk $ fromJust $ find ((chunkIdFromString "MWHD" ==) . chunkId) iff
                   comp = runGet (getWord32be >> getWord8) $ BL.fromChunks [mwhd]
               in decompressBodyF comp body

malwId :: ChunkId
malwId = chunkIdFromString "MALW"

-- | Class for managing chunkable data. This will make it easiser to
-- | handle marshalling of data, etc.
class Chunkable a where
    chunkify :: a -> T
    chunkify x = Cons (chunkifyId x) $ chunkifyData x
    chunkifyId :: a -> ChunkId
    chunkifyData :: a -> Chunk
    -- How to do this?
    -- unchunkify :: T -> Maybe a
    -- Fails with "Chunks.hs:24:51: Not in scope: `a'":
    -- unchunkify (Cons id ch) = if id == chunkableId a then unchunkifyChunk ch else Nothing
    -- unchunkifyChunk :: Chunk -> a
    -- ifOfType :: T -> Bool



instance Chunkable Testtype where
    chunkifyId (Testtype _ _) = chunkIdFromString "TSTY"
    chunkifyData (Testtype w32 w16) = Chunk $ B.concat $ BL.toChunks $ runPut $ do
                                        putWord32be w32
                                        putWord16be w16

makeFORMHISTHIST :: Histogram -> String -> T
makeFORMHISTHIST h s = Cons (chunkIdFromString s) $ Chunk $ B.concat $ BL.toChunks $ runPut $ putF $ elems h
    where putF = foldr ((>>) . putWord32be . fromIntegral) (return ())

makeFORMHIST :: Histogram -> Histogram -> T
makeFORMHIST h r = Cons (chunkIdFromString "HIST") $ Form [makeFORMHISTHIST h "1GRM", makeFORMHISTHIST r "1RGR"]

-- Format for histogram saving:
--
-- FORM.MALW/FORM.HIST: Container for all histogram related stuff
-- FORM.MALW/FORM.HIST/1GRM: normal 1-gram histogram (1D)
-- FORM.MALW/FORM.HIST/1RGR: run-length reduced 1-gram histogram (1D)

-- | Load the histograms from the IFF file. First the normal, second
-- | the rle reduced histogram is returned.
loadFORMHIST :: IFFMap -> (Histogram, Histogram)
loadFORMHIST iff = (lA chunk1grm, lA chunk1rgr)
    where chks = parts $ fromJust $ M.lookup "HIST" iff
          find' n = BL.fromChunks [content $ chunk $ fromJust $ find ((==) (chunkIdFromString n) . chunkId) chks]
          chunk1grm = find' "1GRM"
          chunk1rgr = find' "1RGR"
          lA = listArray (0, 255) . map fromIntegral . runGet (getF 256)
          getF :: Int -> Get [Word32]
          getF 0 = return []
          getF x = getWord32be >>= (\y -> (y :) `fmap` getF (x - 1))


-- PE Headers are just the PE header in binary representation
-- | save a PEHD chunk from a PEHeader
savePEHD :: PEHeader -> T
savePEHD pe = Cons (chunkIdFromString "PEHD") $ Chunk $ B.concat $ BL.toChunks $ runPut putF
  where peoh = peOptionalHeader pe
        putF = do  
          putWord32be $ peSignature pe
          putWord16be $ peMachine pe
          putWord16be $ peNumberOfSections pe
          putWord32be $ peTimeDateStamp pe
          putWord32be $ pePointerToSymbolTable pe
          putWord32be $ peNumberOfSymbols pe
          putWord16be $ peSizeOfOptionalHeader pe
          putWord16be $ peCharacteristics pe
          -- PE optional header from here on...
          putWord16be $ peohMagic peoh
          putWord8 $ peohMajorLinkerVersion peoh
          putWord8 $ peohMinorLinkerVersion peoh
          putWord32be $ peohSizeOfCode peoh
          putWord32be $ peohSizeofInitialisedData peoh
          putWord32be $ peohSizeofUninitialisedData peoh
          putWord32be $ peohAddressOfEntryPoint peoh
          putWord32be $ peohBaseOfCode peoh
          putWord32be $ peohBaseOfData peoh
          putWord32be $ peohImageBase peoh
          putWord32be $ peohSectionAlignment peoh
          putWord32be $ peohFileAlignment peoh
          putWord16be $ peohMajorOSVersion peoh
          putWord16be $ peohMinorOSVersion peoh
          putWord16be $ peohMajorImageVersion peoh
          putWord16be $ peohMinorImageVersion peoh
          putWord16be $ peohMajorSubsystemVersion peoh
          putWord16be $ peohMinorSubsystemVersion peoh
          putWord32be $ peohReserved peoh
          putWord32be $ peohSizeOfImage peoh
          putWord32be $ peohSizeOfHeaders peoh
          putWord32be $ peohCheckSum peoh

-- | load a PEHeader from a Chunk
loadPEHD :: IFFMap -> PEHeader
loadPEHD iff = runGet load chk
    where chk = BL.fromChunks [content $ fromJust $ M.lookup "PEHD" iff]
          load = return PEHeader `ap` getWord32be
                                 `ap` getWord16be
                                 `ap` getWord16be
                                 `ap` getWord32be
                                 `ap` getWord32be
                                 `ap` getWord32be
                                 `ap` getWord16be
                                 `ap` getWord16be
                                 `ap` (PEOptionalHeader `fmap` getWord16be
                                       `ap` getWord8
                                       `ap` getWord8
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be)

-- | Annotation for a file.

-- According to the IFF FORM and Chunk registry an annotation chunk is
-- a generic chunk used for file annotations. We use this one here too
-- (and it will be put into the report).
loadANNO :: IFFMap -> String
loadANNO = f . M.lookup "ANNO"
    where f Nothing = []
          f (Just x) = map (chr . fromIntegral) $ B.unpack $ content x

-- | Annotation chunk (IFF standard)
saveANNO :: String -> T
saveANNO = Cons (chunkIdFromString "ANNO") . Chunk . B.pack . map (fromIntegral . ord)


-- | Save the two standard hashes (currently MD5 and SHA1) to an IFF chunk.
saveHASH :: [Word8] -> SHA1.Word160 -> T
saveHASH m (SHA1.Word160 x1 x2 x3 x4 x5) =  Cons (chunkIdFromString "HASH") $ Form [mkm, mks]
  where mkm = Cons (chunkIdFromString "MD5 ") $ Chunk $ B.pack m
        mks = Cons (chunkIdFromString "SHA1") $ Chunk $ B.concat $ BL.toChunks $ runPut pf
        pf = do putWord32be x1
                putWord32be x2
                putWord32be x3
                putWord32be x4
                putWord32be x5

writeIFF :: String -> [T] -> IO ()
writeIFF fn = B.writeFile fn . toByteString . Cons (chunkIdFromString "MALW") . Form


-- 
